import torch
from time import time
from IPython import display
from matplotlib import pyplot as plt
import numpy as np
import random

""" 通过生成第二特征feature[:,1]和标签labels的散点图，可以直观的感受到两者的线性关系"""
def use_svg_display():
    # 用矢量图表示
    display.set_matplotlib_formats('svg')


def set_figsize(figsize=(3.5, 2.5)):
    use_svg_display()
    # 设置图的尺寸
    plt.rcParams['figure.figsize'] = figsize


# 读取数据
"""在训练模型的时候，我们需要遍历数据集并不断读取⼩批量数据样本。这⾥我们定义⼀个函数：它每次
返回 batch_size （批量⼤⼩）个随机样本的特征和标签。"""

def data_iter(batch_size, features, labels):
    num_examples = len(features) # len(features)=1000
    indices = list(range(num_examples)) # indices = [1 ..... 999]
    random.shuffle(indices)   # 使indices数据随机
    # print(indices)
    for i in range(0, num_examples, batch_size): # 每10个样本取个值
        j = torch.LongTensor(indices[i: min(1+batch_size, num_examples)]) # 创建long类型的tensor
        yield features.index_select(0, j), labels.index_select(0, j) # yield为返回值
"""初识yield表达式是在生成器学习中，只要函数中出现yield表达式，那么这个函数就是生成器函数。当函数执行到yield那一行时，yield就会返回后面的值，然后冻结当前函数状态，下行代码不执行，等下一次调用再唤醒接着往下执行。"""

# 定义模型
def linreg(X, w, b):
    return torch.mm(X, w) + b

# 定义损失函数
def squared_loss(y_hat, y):
    # 注意这⾥返回的是向量
    return (y_hat - y.view(y_hat.size()))**2/2

# 定义优化算法
def sgd(params, lr, batch_size):
    for param in params:
        param.data -= lr * param.grad/batch_size
